package com.sparrow.common.ratelimit;

import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

/**
 * 系统按固定速率生成令牌放入桶中。
 * 每个请求到来时需要消耗一个令牌。
 * 如果令牌足够，请求立即执行；如果没有令牌，请求被限流或阻塞。
 * 可以平滑也可以支持短时突发。
 * 限流算法：令牌桶算法，控制请求处理速率，以保护系统在高并发下的稳定运行
 * 它基于“令牌生产器 + 消费者机制”：
 * 	系统以固定速率向桶中添加“令牌”
 * 	请求到来时，必须“拿到令牌”才被允许执行
 * 	没有令牌 → 拒绝 or 排队 or 等待
 */
public class TokenBucket {
    private final int capacity;     // 桶容量
    private final int rate;         // 每秒令牌数
    private double tokens;          // 当前令牌数量
    private final ScheduledExecutorService scheduler;

    public TokenBucket(int capacity, int rate) {
        if (capacity <= 0 || rate <= 0) {
            throw new IllegalArgumentException("capacity and rate must be positive");
        }

        this.capacity = capacity;
        this.rate = rate;
        this.tokens = capacity;

        // 定时补充令牌
        scheduler = Executors.newSingleThreadScheduledExecutor();
        scheduler.scheduleAtFixedRate(this::addTokens, 0, 100, TimeUnit.MILLISECONDS); // 每 100ms 补充
    }

    // 定时调用：补充令牌
    private synchronized void addTokens() {
        double tokensToAdd = (rate / 10.0);
        if (tokens >= capacity) return;

        double oldTokens = tokens;
        tokens = Math.min(capacity, tokens + tokensToAdd);

        // 只有当整数部分增加时，才唤醒等待线程
        if ((int) tokens > (int) oldTokens) {
            this.notifyAll();
        }
    }

    /**
     * 阻塞式获取 n 个令牌（直到成功为止）
     */
    public synchronized void acquire(int n) throws InterruptedException {
        if (n <= 0) return;
        while (tokens < n) {
            this.wait();
        }
        tokens -= n;
    }

    /**
     * 尝试获取 n 个令牌（不阻塞）
     */
    public synchronized boolean tryAcquire(int n) {
        if (n <= 0) return false;

        if (tokens >= n) {
            tokens -= n;
            return true;
        }
        return false;
    }

    public synchronized double getAvailableTokens() {
        return tokens;
    }

    public void shutdown() {
        scheduler.shutdown();
    }
}

